/*
Copyright 2007 John Tsiombikas <nuclear@siggraph.org>

This file is part of the pixelshow 2007 invitation demo.

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with the program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "kdtree.h"

/* these two definitions are equivalent */
enum { AXIS_X, AXIS_Y, AXIS_Z };
enum { PLANE_YZ, PLANE_XZ, PLANE_XY };

struct kdnode {
	float x, y, z;
	int dir;
	void *data;

	struct kdnode *left, *right;	/* positive/negative side */
};

struct res_node {
	struct kdnode *item;
	float dist_sq;
	struct res_node *next;
};

struct result_set {
	struct res_node *rlist, *riter;
	int size;
};

struct kdtree {
	struct kdnode *root;
	void (*destr)(void*);
};

#define KDTREE(ptr)		((struct kdtree*)ptr)
#define SQ(x)			((x) * (x))


static void clear_rec(struct kdnode *node, void (*destr)(void*));
static int insert_rec(struct kdnode **node, float x, float y, float z, void *data, int dir);
static int rlist_insert(struct res_node *list, struct kdnode *item, float dist_sq);
static void clear_results(struct result_set *tree);
static float axis_value(float x, float y, float z, int dir);

#ifdef USE_LIST_NODE_ALLOCATOR
static void *alloc_resnode(void);
static void free_resnode(void*);
#else
#define alloc_resnode()		malloc(sizeof(struct res_node))
#define free_resnode(n)		free(n)
#endif



void *kd_create(void)
{
	struct kdtree *tree;

	if(!(tree = malloc(sizeof *tree))) {
		return 0;
	}

	tree->root = 0;
	tree->destr = 0;
	return tree;
}

void kd_free(void *tree)
{
	if(tree) {
		kd_clear(tree);
		free(tree);
	}
}

static void clear_rec(struct kdnode *node, void (*destr)(void*))
{
	if(!node) return;

	clear_rec(node->left, destr);
	clear_rec(node->right, destr);
	
	if(destr) {
		destr(node->data);
	}
	free(node);
}

void kd_clear(void *tree)
{
	struct kdtree *kd = KDTREE(tree);
	clear_rec(kd->root, kd->destr);
	kd->root = 0;
}

void kd_data_destructor(void *tree, void (*destr)(void*))
{
	KDTREE(tree)->destr = destr;
}


static int insert_rec(struct kdnode **nptr, float x, float y, float z, void *data, int dir)
{
	int new_dir;
	struct kdnode *node;

	if(!*nptr) {
		if(!(node = malloc(sizeof *node))) {
			return -1;
		}
		node->x = x;
		node->y = y;
		node->z = z;
		node->data = data;
		node->dir = dir;
		node->left = node->right = 0;
		*nptr = node;
		return 0;
	}

	node = *nptr;
	new_dir = (node->dir + 1) % 3;
	if(axis_value(x, y, z, node->dir) < axis_value(node->x, node->y, node->z, node->dir)) {
		return insert_rec(&(*nptr)->left, x, y, z, data, new_dir);
	}
	return insert_rec(&(*nptr)->right, x, y, z, data, new_dir);
}

int kd_insert(void *tree, float x, float y, float z, void *data)
{
	return insert_rec(&KDTREE(tree)->root, x, y, z, data, AXIS_X);
}

/*
void *kd_nearest_n(void *tree, float x, float y, float z, int num)
{
}
*/

static int find_nearest(struct kdnode *node, float x, float y, float z, float range, struct res_node *list, int ordered)
{
	float dist_sq, dx;
	int ret, added_res = 0;

	if(!node) return 0;

	dist_sq = SQ(node->x - x) + SQ(node->y - y) + SQ(node->z - z);
	if(dist_sq <= SQ(range)) {
		if(rlist_insert(list, node, ordered ? dist_sq : -1.0) == -1) {
			return -1;
		}
		added_res = 1;
	}

	dx = axis_value(x, y, z, node->dir) - axis_value(node->x, node->y, node->z, node->dir);

	ret = find_nearest(dx <= 0.0 ? node->left : node->right, x, y, z, range, list, ordered);
	if(ret >= 0 && fabs(dx) < range) {
		added_res += ret;
		ret = find_nearest(dx <= 0.0 ? node->right : node->left, x, y, z, range, list, ordered);
	}
	if(ret == -1) {
		return -1;
	}
	added_res += ret;

	return added_res;
}

void *kd_nearest_range(void *tree, float x, float y, float z, float range)
{
	int ret;
	struct result_set *rset;
	struct kdtree *kd = KDTREE(tree);

	if(!(rset = malloc(sizeof *rset))) {
		return 0;
	}
	if(!(rset->rlist = alloc_resnode())) {
		free(rset);
		return 0;
	}
	rset->rlist->next = 0;

	if((ret = find_nearest(kd->root, x, y, z, range, rset->rlist, 0)) == -1) {
		kd_res_free(rset);
		return 0;
	}
	rset->size = ret;
	kd_res_rewind(rset);
	return rset;
}

void kd_res_free(void *set)
{
	struct result_set *rset = (struct result_set*)set;
	clear_results(set);
	free_resnode(rset->rlist);
	free(rset);
}

int kd_res_size(void *set)
{
	return ((struct result_set*)set)->size;
}

void kd_res_rewind(void *set)
{
	struct result_set *rset = (struct result_set*)set;
	rset->riter = rset->rlist->next;
}

int kd_res_end(void *set)
{
	struct result_set *rset = (struct result_set*)set;
	return rset->riter != 0;
}

int kd_res_next(void *set)
{
	struct result_set *rset = (struct result_set*)set;
	
	rset->riter = rset->riter->next;
	return rset->riter != 0;
}

void *kd_res_item(void *set, float *x, float *y, float *z)
{
	struct result_set *rset = (struct result_set*)set;

	if(rset->riter) {
		if(x) *x = rset->riter->item->x;
		if(y) *y = rset->riter->item->y;
		if(z) *z = rset->riter->item->z;
		return rset->riter->item->data;
	}
	return 0;
}

void *kd_res_item_data(void *set)
{
	return kd_res_item(set, 0, 0, 0);
}

/* ---- static helpers ---- */

#ifdef USE_LIST_NODE_ALLOCATOR
/* special list node allocators.
 * TODO: make thread-safety with pthread_mutex a compile-time option
 */
static struct res_node *free_nodes;

static struct res_node *alloc_resnode(void)
{
	struct res_node *node;

	if(!free_nodes) {
		node = malloc(sizeof *node);
	} else {
		node = free_nodes;
		free_nodes = free_nodes->next;
		node->next = 0;
	}
	return node;
}

static void free_resnode(struct res_node *node)
{
	node->next = free_nodes;
	free_nodes = node;
}
#endif	/* list node allocator or not */


/* inserts the item. if dist_sq is >= 0, then do an ordered insert */
static int rlist_insert(struct res_node *list, struct kdnode *item, float dist_sq)
{
	struct res_node *rnode;

	if(!(rnode = alloc_resnode())) {
		return -1;
	}
	rnode->item = item;
	rnode->dist_sq = dist_sq;

	if(dist_sq >= 0.0) {
		while(list->next && list->next->dist_sq > dist_sq) {
			list = list->next;
		}
	}
	rnode->next = list->next;
	list->next = rnode;
	return 0;
}

static void clear_results(struct result_set *rset)
{
	struct res_node *tmp, *node = rset->rlist->next;

	while(node) {
		tmp = node;
		node = node->next;
		free_resnode(tmp);
	}

	rset->rlist->next = 0;
}

static float axis_value(float x, float y, float z, int dir)
{
	switch(dir) {
	case AXIS_X:
		return x;
	case AXIS_Y:
		return y;
	case AXIS_Z:
		return z;
	default:
		break;
	}
	return 0.0f;
}
